package net.cyue.jpa;

import jakarta.persistence.*;
import jakarta.persistence.spi.PersistenceUnitInfo;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;

public class EntityManagerPool {
    private static final int DEFAULT_SIZE = 20;
    private EntityManagerFactory factory;

    private EntityManagerWrapper[] pool;
    public EntityManagerPool(CYuePersistenceUnitInfoConfiguration configuration) {
        this(configuration, EntityManagerPool.DEFAULT_SIZE);
    }
    public EntityManagerPool(CYuePersistenceUnitInfoConfiguration configuration, int size) {
        PersistenceUnitInfo unitInfo = new CYuePersistenceUnitInfo(configuration);
        this.factory = CYuePersistence.createEntityManagerFactory(unitInfo);
        this.init(size);
    }

    private void init(int size) {
        this.setSize(size);
    }
    public void setSize(int size) {
        if (size < 0) {
            throw new IllegalArgumentException("size must be greater than 0");
        }

        int nowSize = this.pool != null ? this.pool.length : 0;
        if (nowSize == size) {
            return;
        }
        if (nowSize > size) {
            this.pool = Arrays.copyOfRange(this.pool, 0, size);
            return;
        }

        // size > nowSize
        EntityManagerWrapper[] temp = new EntityManagerWrapper[0];
        if (this.pool != null) {
            temp = this.pool.clone();
        }
        this.pool = new EntityManagerWrapper[size];
        System.arraycopy(temp, 0, this.pool, 0, temp.length);
        for (int i = temp.length; i < size; i++) {
            final var manager = this.factory.createEntityManager();
            this.pool[i] = new EntityManagerWrapper(manager);
        }
    }


    private int getAwaitEntityManagerIndex() {
        for (int i = 0; i < this.pool.length; i++) {
            if (!this.pool[i].isUsing()) {
                return i;
            }
        }
        return -1;
    }
    private EntityManagerWrapper getAwaitEntityManager() {
        int index = this.getAwaitEntityManagerIndex();
        if (index == -1) {
            return null;
        }
        return this.pool[index];
    }
    public void useEntityManager(Consumer<EntityManagerWrapper> work) {
        EntityManagerWrapper wrapper = this.getAwaitEntityManager();
        if (wrapper == null) {
            EntityManager manager = this.factory.createEntityManager();
            wrapper = new EntityManagerWrapper(manager);
            work.accept(wrapper);
            wrapper.close();
        }
        work.accept(wrapper);
    }
    public void transaction(Consumer<EntityManager> work) {
        this.useEntityManager((wrapper) -> {
            wrapper.transaction(work);
        });
    }
    public void transactionWrapper(Consumer<EntityManagerWrapper> work) {
        this.useEntityManager((wrapper) -> {
            wrapper.transactionWrapper(work);
        });
    }

    public void query(Consumer<Query> work, String jpql, Object... parameters) {
        this.useEntityManager((wrapper) -> {
            wrapper.query(work, jpql, parameters);
        });
    }
    public void query(Consumer<Query> work, String jpql, Map<String, ?> parameters) {
        this.useEntityManager((wrapper) -> {
            wrapper.query(work, jpql, parameters);
        });
    }
    public <T> void query(Consumer<Query> work, Class<T> clz, String... columnName) {
        this.useEntityManager((wrapper) -> {
            wrapper.query(work, clz, columnName);
        });
    }


    public <T> List<T> select(Class<T> clz, String... columnName) {
        AtomicReference<List<T>> ref = new AtomicReference<>();
        this.useEntityManager((wrapper) -> {
            ref.set(wrapper.select(clz, columnName));
        });
        return ref.get();
    }
    public <T> T selectOne(Class<T> clz, String... columnName) {
        AtomicReference<T> ref = new AtomicReference<>();
        this.useEntityManager(
            (wrapper) -> {
                ref.set(wrapper.selectOne(clz, columnName));
            }
        );
        return ref.get();
    }

    public Object selectOne(String jpql, Object... parameters) {
        AtomicReference<Object> ref = new AtomicReference<>();
        this.useEntityManager(
            (wrapper) -> {
                ref.set(wrapper.selectOne(jpql, parameters));
            }
        );
        return ref.get();
    }
    public Object selectOne(String jpql, Map<String, ?> parameters) {
        AtomicReference<Object> ref = new AtomicReference<>();
        this.useEntityManager(
            (wrapper) -> {
                ref.set(wrapper.selectOne(jpql, parameters));
            }
        );
        return ref.get();
    }
    public List select(String jpql, Object... parameters) {
        AtomicReference<List<?>> ref = new AtomicReference<>();
        this.useEntityManager(
            (wrapper) -> {
                ref.set(wrapper.select(jpql, parameters));
            }
        );
        return ref.get();
    }
    public List select(String jpql, Map<String, ?> parameters) {
        AtomicReference<List<?>> ref = new AtomicReference<>();
        this.useEntityManager(
            (wrapper) -> {
                ref.set(wrapper.select(jpql, parameters));
            }
        );
        return ref.get();
    }
    public int executeUpdate(String jpql, Object... parameters) {
        AtomicInteger ref = new AtomicInteger();
        this.useEntityManager(
            (wrapper) -> {
                ref.set(wrapper.executeUpdate(jpql, parameters));
            }
        );
        return ref.get();
    }
    public int executeUpdate(String jpql, Map<String, ?> parameters) {
        AtomicInteger ref = new AtomicInteger();
        this.useEntityManager(
            (wrapper) -> {
                ref.set(wrapper.executeUpdate(jpql, parameters));
            }
        );
        return ref.get();
    }


    public void persist(Object obj) {
        this.useEntityManager((wrapper) -> wrapper.persist(obj));
    }

    public <T> T find(Class<T> clz, int id) {
        AtomicReference<T> ref = new AtomicReference<>();
        this.useEntityManager((wrapper) -> {
            ref.set(wrapper.find(clz, id));
        });
        return ref.get();
    }

    public void remove(Object obj) {
        this.useEntityManager((wrapper) -> wrapper.remove(obj));
    }
    public <T> void remove(Class<T> clz, int id) {
        this.useEntityManager((wrapper) -> wrapper.remove(clz, id));
    }
}

